import numpy as np
import open3d as o3d
from scipy.spatial import distance_matrix

from cuboid_locate import cuboid_locate, pass_through_filter


def compute_nearest_distance(pcd: o3d.geometry.PointCloud) -> float:
    points = np.asarray(pcd.points)
    distance = distance_matrix(points, points)
    sorted_distance = np.sort(distance)
    nearest_distance = sorted_distance[:, 1]
    mean = np.mean(nearest_distance)
    std = np.mean(nearest_distance)
    return mean + 3 * std


def main():
    # 第2组
    # pcd_file_path = R"resource/stack/第二组/20220417T103403_1201_1.pcd"
    # pcd_file_path = R"resource/stack/第二组/20220417T103310_1201_1.pcd"
    # pcd_file_path = R"resource/stack/第二组/20220417T103234_1201_1.pcd"
    # 第3组
    pcd_file_path = R"resource/stack/第三组/20220417T104700_1201_1.pcd"
    # pcd_file_path = R"resource/stack/第三组/20220417T104617_1201_1.pcd"
    # pcd_file_path = R"resource/stack/第三组/20220417T104550_1201_1.pcd"
    # 第5组
    # pcd_file_path = R"resource/stack/第五组/20220417T111324_1201_1.pcd"
    # pcd_file_path = R"resource/stack/第五组/20220417T111237_1201_1.pcd"
    # pcd_file_path = R"resource/stack/第五组/20220417T111214_1201_1.pcd"
    # 第6组
    # pcd_file_path = R"resource/stack/第六组/20220417T112307_1201_1.pcd"
    # pcd_file_path = R"resource/stack/第六组/20220417T112232_1201_1.pcd"
    # pcd_file_path = R"resource/stack/第六组/20220417T112140_1201_1.pcd"

    x_range = (-400, 160)
    y_range = (-100, 200)
    z_range = (700, 800)
    spot_distance = 10
    cuboid_length = 305
    cuboid_width = 120

    pcd = o3d.io.read_point_cloud(pcd_file_path)
    pcd, surface = cuboid_locate(pcd, x_range, y_range, z_range, spot_distance,
                                 cuboid_length, cuboid_width)

    top_cuboid_surface_height = surface.top_left.z
    z_range = (top_cuboid_surface_height + 100, 1500)
    origin_pcd = o3d.io.read_point_cloud(pcd_file_path)
    pcd = pass_through_filter(origin_pcd, x_range, y_range, z_range)
    ground_height = np.median(np.asarray(pcd.points)[:, -1])

    x_range = (surface.top_left.x, surface.top_right.x)
    y_low = min(surface.top_left.y, surface.top_right.y)
    y_high = max(surface.top_left.y, surface.top_right.y)
    y_range = (y_low - 50, y_high + 50)
    z_range = (top_cuboid_surface_height + 50, ground_height - 50)
    origin_pcd = o3d.io.read_point_cloud(pcd_file_path)
    top_pcd = pass_through_filter(origin_pcd, x_range, y_range, z_range)
    radius = compute_nearest_distance(top_pcd)
    top_pcd, _ = top_pcd.remove_radius_outlier(nb_points=2, radius=radius)
    top_pcd.paint_uniform_color([1, 0, 0])

    x_range = (surface.bottom_left.x, surface.bottom_right.x)
    y_low = min(surface.bottom_left.y, surface.bottom_right.y)
    y_high = max(surface.bottom_left.y, surface.bottom_right.y)
    y_range = (y_low - 50, y_high + 50)
    z_range = (top_cuboid_surface_height + 50, ground_height - 50)
    origin_pcd = o3d.io.read_point_cloud(pcd_file_path)
    bottom_pcd = pass_through_filter(origin_pcd, x_range, y_range, z_range)
    radius = compute_nearest_distance(bottom_pcd)
    bottom_pcd, _ = bottom_pcd.remove_radius_outlier(nb_points=2, radius=radius)
    bottom_pcd.paint_uniform_color([0, 0, 1])

    bottom_points = np.asarray(bottom_pcd.points)
    with open(R"E:\temp\bottom_points.npy", "wb") as f:
        np.save(f, bottom_points)

    # top_mean_z = np.mean(np.asarray(top_pcd.points)[:, -1])
    # bottom_mean_z = np.mean(np.asarray(bottom_pcd.points)[:, -1])
    # abs_z_diff = abs(top_mean_z - bottom_mean_z)
    # print(F"abs z diff: {abs_z_diff}")

    o3d.visualization.draw_geometries([top_pcd, bottom_pcd])


if __name__ == "__main__":
    main()
